#!/usr/bin/env python 
# -*- coding:utf-8 -*-

from sklearn.naive_bayes import GaussianNB
import numpy as np
import sys


def get_data(path):
    train = []
    with open(path, "r", encoding="utf8") as file:
        for line in file:
            train.append(line.strip().split(","))
    test = np.asarray(train.pop(), dtype=np.float32)
    train = np.asarray(train, dtype=np.float32)
    return train, test


def get_y_train_data(path):
    y = {}
    with open(path, "r", encoding="utf8") as file:
        for line in file:
            line_list = line.strip().split(",")
            station_id = line_list.pop(0)
            y[station_id] = np.asarray(line_list, dtype=np.float32)
    return y


if __name__ == '__main__':
    base_path = ""                    #sys.argv[1]
    x_file_name = "X1565160402221.txt"#sys.argv[2]
    y_file_name = "Y1565160402221.txt"#sys.argv[3]
    predict_file_name = "value.txt"   #sys.argv[4]
    # train_num = 6000
    # (x_train, y_train), (x_test, y_test) = get_dataset(
    #     "mushroom.txt", n_train=train_num, tar_idx=0)
    # print(np.mean(y_test == clf.predict(x_test_one_hot)))
    # 入参训练数据读取
    (x_train, x_test) = get_data(base_path + x_file_name)
    # 出参训练数据读取
    y_train_dict = get_y_train_data(base_path + y_file_name)
    clf = GaussianNB()
    # 创建预测值文件对象
    predict_file = open(base_path + predict_file_name, "w")
    for (station, y_train) in y_train_dict.items():
        clf.fit(x_train, y_train)
        y_predict = clf.predict(x_test.reshape(1, -1))
        # 写预测值
        predict_file.write(station + "," + str(y_predict[0]) + "\n")
    predict_file.close()
